#include "loop_closure/Scancontext.h"
#include <algorithm> // For std::sort

// A helper function not in the class
float xy2theta(const float &_x, const float &_y)
{
    if (_x >= 0 && _y >= 0)
        return (180.0 / M_PI) * atan(_y / _x);
    if (_x < 0 && _y >= 0)
        return 180.0 - ((180.0 / M_PI) * atan(_y / (-_x)));
    if (_x < 0 && _y < 0)
        return 180.0 + ((180.0 / M_PI) * atan(_y / _x));
    if (_x >= 0 && _y < 0)
        return 360.0 - ((180.0 / M_PI) * atan((-_y) / _x));
    return 0; // Should not happen
}

// A helper function not in the class
MatrixXd circshift(MatrixXd &_mat, int _num_shift)
{
    assert(_num_shift >= 0);
    if (_num_shift == 0)
        return _mat;

    MatrixXd shifted_mat = MatrixXd::Zero(_mat.rows(), _mat.cols());
    for (int col_idx = 0; col_idx < _mat.cols(); col_idx++)
    {
        int new_location = (col_idx + _num_shift) % _mat.cols();
        shifted_mat.col(new_location) = _mat.col(col_idx);
    }
    return shifted_mat;
}

// A helper function not in the class
std::vector<float> eig2stdvec(MatrixXd _eigmat)
{
    std::vector<float> vec(_eigmat.data(), _eigmat.data() + _eigmat.size());
    return vec;
}

// Class constructor
SCManager::SCManager() {}

// Main algorithm implementations
MatrixXd SCManager::makeScancontext(pcl::PointCloud<SCPointType> &_scan_down)
{
    TicToc t_making_desc;

    const int NO_POINT = -1000;
    MatrixXd desc = NO_POINT * MatrixXd::Ones(PC_NUM_RING, PC_NUM_SECTOR);

    for (const auto &pt : _scan_down.points)
    {
        float azim_range = sqrt(pt.x * pt.x + pt.y * pt.y);
        if (azim_range > PC_MAX_RADIUS)
            continue;

        float azim_angle = xy2theta(pt.x, pt.y);

        int ring_idx = std::max(0, std::min(PC_NUM_RING - 1, int(floor(azim_range / (PC_MAX_RADIUS / PC_NUM_RING)))));
        int sctor_idx = std::max(0, std::min(PC_NUM_SECTOR - 1, int(floor(azim_angle / PC_UNIT_SECTORANGLE))));

        float current_z = pt.z + LIDAR_HEIGHT;
        if (desc(ring_idx, sctor_idx) < current_z)
        {
            desc(ring_idx, sctor_idx) = current_z;
        }
    }

    for (int row_idx = 0; row_idx < desc.rows(); row_idx++)
        for (int col_idx = 0; col_idx < desc.cols(); col_idx++)
            if (desc(row_idx, col_idx) == NO_POINT)
                desc(row_idx, col_idx) = 0;

    t_making_desc.toc("Scan Context making");
    return desc;
}

MatrixXd SCManager::makeRingkeyFromScancontext(Eigen::MatrixXd &_desc)
{
    Eigen::MatrixXd invariant_key(_desc.rows(), 1);
    for (int row_idx = 0; row_idx < _desc.rows(); row_idx++)
    {
        invariant_key(row_idx, 0) = _desc.row(row_idx).mean();
    }
    return invariant_key;
}

MatrixXd SCManager::makeSectorkeyFromScancontext(Eigen::MatrixXd &_desc)
{
    Eigen::MatrixXd variant_key(1, _desc.cols());
    for (int col_idx = 0; col_idx < _desc.cols(); col_idx++)
    {
        variant_key(0, col_idx) = _desc.col(col_idx).mean();
    }
    return variant_key;
}

void SCManager::makeAndSaveScancontextAndKeys(pcl::PointCloud<SCPointType> &_scan_down)
{
    Eigen::MatrixXd sc = makeScancontext(_scan_down);
    Eigen::MatrixXd ringkey = makeRingkeyFromScancontext(sc);
    Eigen::MatrixXd sectorkey = makeSectorkeyFromScancontext(sc);
    std::vector<float> ringkey_stdvec = eig2stdvec(ringkey);

    polarcontexts_.push_back(sc);
    polarcontext_vkeys_.push_back(sectorkey);
    polarcontext_invkeys_mat_.push_back(ringkey_stdvec);
}

double SCManager::distDirectSC(MatrixXd &_sc1, MatrixXd &_sc2)
{
    int num_eff_cols = 0;
    double sum_sector_similarity = 0;
    for (int col_idx = 0; col_idx < _sc1.cols(); col_idx++)
    {
        VectorXd col_sc1 = _sc1.col(col_idx);
        VectorXd col_sc2 = _sc2.col(col_idx);

        if (col_sc1.norm() == 0 || col_sc2.norm() == 0)
            continue;

        double sector_similarity = col_sc1.dot(col_sc2) / (col_sc1.norm() * col_sc2.norm());
        sum_sector_similarity += sector_similarity;
        num_eff_cols++;
    }

    if (num_eff_cols == 0)
        return 1.0; // No overlap, max distance

    double sc_sim = sum_sector_similarity / num_eff_cols;
    return 1.0 - sc_sim;
}

int SCManager::fastAlignUsingVkey(MatrixXd &_vkey1, MatrixXd &_vkey2)
{
    int argmin_vkey_shift = 0;
    double min_veky_diff_norm = std::numeric_limits<double>::max();
    for (int shift_idx = 0; shift_idx < _vkey1.cols(); shift_idx++)
    {
        MatrixXd vkey2_shifted = circshift(_vkey2, shift_idx);
        double cur_diff_norm = (_vkey1 - vkey2_shifted).norm();
        if (cur_diff_norm < min_veky_diff_norm)
        {
            argmin_vkey_shift = shift_idx;
            min_veky_diff_norm = cur_diff_norm;
        }
    }
    return argmin_vkey_shift;
}

std::pair<double, int> SCManager::distanceBtnScanContext(MatrixXd &_sc1, MatrixXd &_sc2)
{
    MatrixXd vkey_sc1 = makeSectorkeyFromScancontext(_sc1);
    MatrixXd vkey_sc2 = makeSectorkeyFromScancontext(_sc2);
    int argmin_vkey_shift = fastAlignUsingVkey(vkey_sc1, vkey_sc2);

    const int SEARCH_RADIUS = round(0.5 * 0.1 * _sc1.cols()); // Search ratio 10%
    std::vector<int> shift_idx_search_space;
    for (int i = -SEARCH_RADIUS; i <= SEARCH_RADIUS; ++i)
    {
        shift_idx_search_space.push_back((argmin_vkey_shift + i + _sc1.cols()) % _sc1.cols());
    }

    int argmin_shift = 0;
    double min_sc_dist = std::numeric_limits<double>::max();
    for (int num_shift : shift_idx_search_space)
    {
        MatrixXd sc2_shifted = circshift(_sc2, num_shift);
        double cur_sc_dist = distDirectSC(_sc1, sc2_shifted);
        if (cur_sc_dist < min_sc_dist)
        {
            argmin_shift = num_shift;
            min_sc_dist = cur_sc_dist;
        }
    }
    return make_pair(min_sc_dist, argmin_shift);
}

std::pair<int, float> SCManager::detectLoopClosureID(void)
{
    int loop_id = -1; 
    
    if ((int)polarcontext_invkeys_mat_.size() < NUM_EXCLUDE_RECENT + 1)
    {
        return make_pair(-1, 0.0);
    }

    if (tree_making_period_conter % TREE_MAKING_PERIOD_ == 0)
    {
        TicToc t_tree_construction;
        
        // ======================== 核心修正-3: 修正KD-Tree的创建 ========================
        // 创建适配器实例
        KeyMat_Adaptor adaptor(polarcontext_invkeys_mat_);
        // 使用适配器创建KD-Tree
        polarcontext_tree_ = std::make_unique<InvKeyTree>(PC_NUM_RING, adaptor, nanoflann::KDTreeSingleIndexAdaptorParams(10));
        // ==============================================================================
        
        t_tree_construction.toc("Tree construction");
    }
    tree_making_period_conter++;

    auto& curr_key = polarcontext_invkeys_mat_.back();
    auto& curr_desc = polarcontexts_.back();

    std::vector<size_t> ret_indexes(NUM_CANDIDATES_FROM_TREE);
    std::vector<float> out_dists_sqr(NUM_CANDIDATES_FROM_TREE);

    nanoflann::KNNResultSet<float> resultSet(NUM_CANDIDATES_FROM_TREE);
    resultSet.init(&ret_indexes[0], &out_dists_sqr[0]);
    
    // ======================== 核心修正-4: 修正findNeighbors的调用 ========================
    // 直接在 unique_ptr 上调用 findNeighbors
    polarcontext_tree_->findNeighbors(resultSet, &curr_key[0], nanoflann::SearchParameters(10));
    // ==============================================================================

    double min_dist = std::numeric_limits<double>::max();
    int nn_align = 0;
    int nn_idx = -1;

    for (size_t i = 0; i < resultSet.size(); ++i)
    {
        int candidate_idx = ret_indexes[i];
        
        if (abs((int)polarcontexts_.size() - 1 - candidate_idx) < NUM_EXCLUDE_RECENT)
            continue;

        MatrixXd& candidate_desc = polarcontexts_[candidate_idx];
        auto sc_dist_result = distanceBtnScanContext(curr_desc, candidate_desc);
        
        if (sc_dist_result.first < min_dist)
        {
            min_dist = sc_dist_result.first;
            nn_align = sc_dist_result.second;
            nn_idx = candidate_idx;
        }
    }

    if (min_dist < SC_DIST_THRES)
    {
        loop_id = nn_idx;
    }

    float yaw_diff_rad = (float)nn_align * PC_UNIT_SECTORANGLE * M_PI / 180.0;
    return make_pair(loop_id, yaw_diff_rad);
}